"""
RedFire框架命名数据库连接管理器

提供企业级的多数据库连接管理、事务处理和智能路由。
支持：
- 多命名数据库连接
- 读写分离自动路由
- 异构数据库技术栈
- 负载均衡和故障转移
- 强制服务隔离
"""

from typing import Dict, Any, Optional, AsyncGenerator, List, Set, Union
import logging
from contextlib import asynccontextmanager
from enum import Enum
import asyncio

try:
    from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, AsyncEngine
    from sqlalchemy.orm import sessionmaker, DeclarativeBase
    from sqlalchemy.pool import StaticPool, QueuePool
    from sqlalchemy import text
    
    class Base(DeclarativeBase):
        """SQLAlchemy声明式基类"""
        pass
        
except ImportError:
    # 如果没有安装SQLAlchemy，提供模拟实现
    AsyncSession = None
    AsyncEngine = None
    create_async_engine = None
    sessionmaker = None
    StaticPool = None
    QueuePool = None
    text = None
    
    class Base:
        """模拟SQLAlchemy基类"""
        __tablename__ = ""
        pass

try:
    import redis.asyncio as redis
    from redis.asyncio import Redis
except ImportError:
    # 如果没有安装redis，提供模拟实现
    redis = None
    Redis = None

logger = logging.getLogger(__name__)


class DatabaseRole(Enum):
    """数据库角色枚举"""
    WRITE = "write"      # 写入角色（主库）
    READ = "read"        # 读取角色（只读副本）
    CACHE = "cache"      # 缓存角色（Redis等）
    ANALYTICS = "analytics"  # 分析角色（时序数据库等）


class DatabaseType(Enum):
    """数据库类型枚举"""
    POSTGRESQL = "postgresql"
    MYSQL = "mysql"
    SQLITE = "sqlite"
    REDIS = "redis"
    INFLUXDB = "influxdb"
    MONGODB = "mongodb"


class DatabaseManager:
    """
    增强型数据库连接管理器
    
    企业级多数据库连接管理器，支持：
    - 多个命名数据库连接
    - 读写分离和负载均衡
    - 异构数据库技术栈
    - 自动化生命周期管理
    - 向后兼容的配置格式
    - 智能连接路由
    """
    
    def __init__(self, databases_config: Optional[Union[Dict[str, Dict[str, Any]], Dict[str, Any]]] = None):
        """
        初始化数据库管理器
        
        Args:
            databases_config: 数据库配置，支持两种格式：
            
            1. 嵌套格式（多数据库）:
            {
                "user_primary": {
                    "url": "postgresql+asyncpg://host/user_db",
                    "pool_size": 5,
                    "role": "write",
                    "type": "postgresql"
                },
                "user_readonly": {
                    "url": "postgresql+asyncpg://readonly-host/user_db", 
                    "pool_size": 20,
                    "role": "read",
                    "type": "postgresql"
                }
            }
            
            2. 扁平格式（单数据库，向后兼容）:
            {
                "url": "postgresql+asyncpg://host/db",
                "pool_size": 10,
                "echo": False
            }
        """
        self.engines: Dict[str, AsyncEngine] = {}
        self.session_factories: Dict[str, sessionmaker] = {}
        self._initialized_databases: Set[str] = set()
        
        # 处理配置格式兼容性
        if databases_config is None:
            self.databases_config = {}
        else:
            self.databases_config = self._normalize_config(databases_config)
        
        # 验证并自动初始化所有数据库连接
        self._validate_config()
        self._initialize_all_sync()
    
    def _normalize_config(self, config: Union[Dict[str, Dict[str, Any]], Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
        """
        标准化配置格式，支持向后兼容
        
        Args:
            config: 原始配置（扁平或嵌套格式）
            
        Returns:
            标准化的嵌套格式配置
        """
        # 检查是否是扁平格式（包含url字段）
        if "url" in config and isinstance(config.get("url"), str):
            # 扁平格式 -> 转换为嵌套格式，使用"primary"作为默认名称
            logger.info("检测到扁平格式数据库配置，自动转换为嵌套格式")
            return {
                "primary": dict(config)  # 复制配置避免修改原始数据
            }
        
        # 已经是嵌套格式，直接返回
        return config
    
    def _validate_config(self) -> None:
        """验证数据库配置的有效性"""
        if not self.databases_config:
            logger.info("没有数据库配置，创建空的数据库管理器")
            return
        
        for db_name, config in self.databases_config.items():
            if not isinstance(config, dict):
                raise ValueError(f"数据库 '{db_name}' 的配置必须是字典格式")
            
            if "url" not in config:
                raise ValueError(f"数据库 '{db_name}' 缺少必需的 'url' 配置")
            
            # 设置默认值
            config.setdefault("pool_size", 10)
            config.setdefault("pool_timeout", 30)
            config.setdefault("echo", False)
            config.setdefault("role", "write")
            config.setdefault("type", self._detect_db_type(config["url"]))
    
    def _detect_db_type(self, url: str) -> str:
        """根据URL自动检测数据库类型"""
        url_lower = url.lower()
        if "postgresql" in url_lower or "postgres" in url_lower:
            return DatabaseType.POSTGRESQL.value
        elif "mysql" in url_lower:
            return DatabaseType.MYSQL.value
        elif "sqlite" in url_lower:
            return DatabaseType.SQLITE.value
        elif "redis" in url_lower:
            return DatabaseType.REDIS.value
        elif "influx" in url_lower:
            return DatabaseType.INFLUXDB.value
        elif "mongo" in url_lower:
            return DatabaseType.MONGODB.value
        else:
            logger.warning(f"无法检测数据库类型，URL: {url}")
            return "unknown"
    
    def _initialize_all_sync(self) -> None:
        """
        同步初始化所有数据库连接
        
        在构造函数中调用，确保实例创建后立即可用
        """
        if not self.databases_config:
            logger.info("没有数据库配置，跳过初始化")
            return
            
        if not create_async_engine:
            logger.warning("SQLAlchemy未安装，使用模拟数据库管理器")
            return
            
        for db_name, config in self.databases_config.items():
            try:
                self._initialize_single_database_sync(db_name, config)
                logger.info(f"数据库 '{db_name}' ({config['type']}) 初始化完成")
            except Exception as e:
                logger.error(f"数据库 '{db_name}' 初始化失败: {e}")
                raise
        
        logger.info(f"所有数据库连接初始化完成，共 {len(self.engines)} 个连接")
    
    def _initialize_single_database_sync(self, db_name: str, config: Dict[str, Any]) -> None:
        """
        同步初始化单个数据库连接
        
        Args:
            db_name: 数据库逻辑名称
            config: 数据库配置
        """
        # 特殊处理Redis数据库
        if config["type"] == DatabaseType.REDIS.value:
            self._initialize_redis_connection(db_name, config)
            return
        
        # 处理关系型数据库（PostgreSQL, MySQL, SQLite）
        # 根据数据库类型选择合适的连接池
        # 注意：异步引擎不能使用QueuePool，需要使用默认的连接池
        poolclass = None
        if config["type"] == DatabaseType.SQLITE.value:
            poolclass = StaticPool
        # 对于PostgreSQL和MySQL的异步引擎，不指定poolclass，让SQLAlchemy使用默认的
        
            # 创建异步引擎
        engine_kwargs = {
            "echo": config["echo"]
        }
        
        # 只在需要时设置poolclass
        if poolclass is not None:
            engine_kwargs["poolclass"] = poolclass
        
        # SQLite不支持pool_size和pool_timeout参数
        if config["type"] != DatabaseType.SQLITE.value:
            engine_kwargs.update({
                "pool_size": config["pool_size"],
                "pool_timeout": config["pool_timeout"]
            })
        
        engine = create_async_engine(config["url"], **engine_kwargs)
        
        # 创建会话工厂（仅对关系型数据库）
        if config["type"] in [DatabaseType.POSTGRESQL.value, DatabaseType.MYSQL.value, DatabaseType.SQLITE.value]:
            session_factory = sessionmaker(
                bind=engine,
                class_=AsyncSession,
                expire_on_commit=False
            )
            self.session_factories[db_name] = session_factory
        
        self.engines[db_name] = engine
        self._initialized_databases.add(db_name)
    
    def _initialize_redis_connection(self, db_name: str, config: Dict[str, Any]) -> None:
        """
        初始化Redis连接
        
        Args:
            db_name: 数据库逻辑名称
            config: Redis配置
        """
        if Redis is None:
            raise ImportError("Redis库未安装，请运行: pip install 'redis[hiredis]'")
        
        # 解析Redis URL
        url = config["url"]
        
        # 创建Redis连接池
        redis_client = redis.from_url(
            url,
            encoding="utf-8",
            decode_responses=True,
            socket_timeout=config.get("pool_timeout", 30),
            socket_connect_timeout=config.get("pool_timeout", 30),
            max_connections=config.get("pool_size", 10)
        )
        
        # 将Redis客户端存储为"引擎"（保持接口一致性）
        self.engines[db_name] = redis_client
        self._initialized_databases.add(db_name)
        
        logger.info(f"数据库 '{db_name}' (redis) 初始化完成")
        
    async def initialize_database(self, db_name: str) -> None:
        """
        按需初始化指定的数据库连接（已废弃，保留向后兼容）
        
        Args:
            db_name: 数据库逻辑名称
        """
        logger.warning("initialize_database方法已废弃，数据库连接在构造函数中自动初始化")
        if db_name not in self.engines:
            raise ValueError(f"数据库 '{db_name}' 未配置或初始化失败")
    
    async def initialize_all(self) -> None:
        """初始化所有配置的数据库连接（已废弃，保留向后兼容）"""
        logger.warning("initialize_all方法已废弃，数据库连接在构造函数中自动初始化")
        if not self.engines:
            raise RuntimeError("没有可用的数据库连接，请检查配置")
    
    async def close_all(self) -> None:
        """关闭所有数据库连接"""
        close_tasks = []
        for db_name, engine in self.engines.items():
            if not engine:
                continue

            # 根据引擎类型选择不同的关闭方法
            # Redis客户端有 aclose 或 close 方法
            if hasattr(engine, 'aclose'):
                # for redis >= 5.0
                close_tasks.append(engine.aclose())
            elif hasattr(engine, 'close') and asyncio.iscoroutinefunction(engine.close):
                 # for redis < 5.0
                close_tasks.append(engine.close())
            elif hasattr(engine, 'dispose'):
                # SQLAlchemy 引擎
                close_tasks.append(engine.dispose())
            else:
                logger.warning(f"无法确定数据库 '{db_name}' 的关闭方法")

        if close_tasks:
            results = await asyncio.gather(*close_tasks, return_exceptions=True)
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    db_name = list(self.engines.keys())[i]
                    logger.error(f"关闭数据库 '{db_name}' 连接时发生错误: {result}")
            
            logger.info(f"已尝试关闭 {len(close_tasks)} 个数据库连接")
        
        self.engines.clear()
        self.session_factories.clear()
        self._initialized_databases.clear()
    
    @asynccontextmanager
    async def get_session(self, db_name: Optional[str] = None) -> AsyncGenerator[AsyncSession, None]:
        """
        获取数据库会话上下文管理器
        
        Args:
            db_name: 数据库逻辑名称，如果为None则使用"primary"（向后兼容）
        
        Yields:
            AsyncSession: 数据库会话
        """
        # 向后兼容：如果没有指定db_name，使用"primary"
        if db_name is None:
            db_name = "primary"
        
        if db_name not in self.session_factories:
            raise ValueError(f"数据库 '{db_name}' 不支持会话操作（可能是非关系型数据库或未配置）")
        
        session_factory = self.session_factories[db_name]
        async with session_factory() as session:
            try:
                yield session
            except Exception as e:
                logger.error(f"数据库 '{db_name}' 会话错误: {e}")
                await session.rollback()
                raise
            finally:
                await session.close()
    
    async def get_engine(self, db_name: Optional[str] = None) -> Union[AsyncEngine, Redis]:
        """
        获取数据库引擎或连接
        
        Args:
            db_name: 数据库逻辑名称，如果为None则使用"primary"（向后兼容）
            
        Returns:
            Union[AsyncEngine, Redis]: 数据库引擎（SQLAlchemy）或Redis客户端
        """
        # 向后兼容：如果没有指定db_name，使用"primary"
        if db_name is None:
            db_name = "primary"
        
        if db_name not in self.engines:
            raise ValueError(f"数据库 '{db_name}' 未配置或初始化失败")
        
        return self.engines[db_name]
    
    async def execute_raw_sql(self, db_name: str, sql: str, params: Optional[Dict[str, Any]] = None) -> Any:
        """
        在指定数据库上执行原生SQL查询
        
        Args:
            db_name: 数据库逻辑名称
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        async with self.get_session(db_name) as session:
            try:
                result = await session.execute(text(sql), params or {})
                await session.commit()
                return result
            except Exception as e:
                logger.error(f"数据库 '{db_name}' SQL执行失败: {e}")
                await session.rollback()
                raise
    
    def list_databases(self) -> List[str]:
        """列出所有配置的数据库名称"""
        return list(self.databases_config.keys())
    
    def has_database(self, db_name: str) -> bool:
        """检查是否存在指定的数据库配置"""
        return db_name in self.databases_config
    
    def get_database_info(self, db_name: str) -> Dict[str, Any]:
        """获取数据库配置信息"""
        if db_name not in self.databases_config:
            raise ValueError(f"未找到数据库配置: {db_name}")
        
        config = self.databases_config[db_name].copy()
        # 隐藏敏感信息
        if "url" in config:
            config["url"] = self._mask_url(config["url"])
        return config
    
    def _mask_url(self, url: str) -> str:
        """遮蔽URL中的敏感信息"""
        import re
        # 遮蔽密码部分
        return re.sub(r'://([^:]+):([^@]+)@', r'://\1:***@', url)
    
    def is_database_initialized(self, db_name: str) -> bool:
        """检查指定数据库是否已初始化"""
        return db_name in self._initialized_databases
    
    def get_databases_by_role(self, role: Union[str, DatabaseRole]) -> List[str]:
        """根据角色获取数据库列表"""
        if isinstance(role, DatabaseRole):
            role = role.value
        
        return [
            db_name for db_name, config in self.databases_config.items()
            if config.get("role") == role
        ]
    
    def get_databases_by_type(self, db_type: Union[str, DatabaseType]) -> List[str]:
        """根据类型获取数据库列表"""
        if isinstance(db_type, DatabaseType):
            db_type = db_type.value
        
        return [
            db_name for db_name, config in self.databases_config.items()
            if config.get("type") == db_type
        ]


class DatabaseRouter:
    """
    数据库连接路由器
    
    提供高级的数据库路由功能：
    - 读写分离自动路由
    - 负载均衡
    - 故障转移
    """
    
    def __init__(self, db_manager: DatabaseManager):
        """
        初始化数据库路由器
        
        Args:
            db_manager: 数据库管理器实例
        """
        self.db_manager = db_manager
        self._read_replicas_index = 0  # 用于负载均衡的索引
    
    @asynccontextmanager
    async def get_write_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        获取写入会话（自动路由到主库）
        
        Yields:
            AsyncSession: 写入数据库会话
        """
        write_databases = self.db_manager.get_databases_by_role(DatabaseRole.WRITE)
        if not write_databases:
            raise ValueError("未找到写入数据库配置")
        
        # 使用第一个写入数据库（通常只有一个主库）
        primary_db = write_databases[0]
        async with self.db_manager.get_session(primary_db) as session:
            yield session
    
    @asynccontextmanager
    async def get_read_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        获取读取会话（自动负载均衡到只读副本）
        
        Yields:
            AsyncSession: 读取数据库会话
        """
        read_databases = self.db_manager.get_databases_by_role(DatabaseRole.READ)
        
        # 如果没有只读副本，回退到写入数据库
        if not read_databases:
            write_databases = self.db_manager.get_databases_by_role(DatabaseRole.WRITE)
            if not write_databases:
                raise ValueError("未找到可用的数据库配置")
            read_databases = write_databases
        
        # 简单的轮询负载均衡
        selected_db = read_databases[self._read_replicas_index % len(read_databases)]
        self._read_replicas_index += 1
        
        async with self.db_manager.get_session(selected_db) as session:
            yield session
    
    @asynccontextmanager
    async def get_session_for_operation(self, operation_type: str) -> AsyncGenerator[AsyncSession, None]:
        """
        根据操作类型获取最优会话
        
        Args:
            operation_type: 操作类型 ("read", "write", "transaction")
            
        Yields:
            AsyncSession: 数据库会话
        """
        if operation_type.lower() in ["write", "insert", "update", "delete", "transaction"]:
            async with self.get_write_session() as session:
                yield session
        else:
            async with self.get_read_session() as session:
                yield session
    
    async def get_cache_session(self, cache_name: Optional[str] = None) -> Any:
        """
        获取缓存会话
        
        Args:
            cache_name: 缓存数据库名称，如果为None则使用第一个缓存数据库
            
        Returns:
            缓存会话对象
        """
        cache_databases = self.db_manager.get_databases_by_role(DatabaseRole.CACHE)
        if not cache_databases:
            raise ValueError("未找到缓存数据库配置")
        
        target_cache = cache_name if cache_name and cache_name in cache_databases else cache_databases[0]
        
        # 对于缓存数据库，直接返回引擎（Redis等不使用SQLAlchemy会话）
        return await self.db_manager.get_engine(target_cache)
    
    def get_database_status(self) -> Dict[str, Any]:
        """获取所有数据库的状态信息"""
        status = {
            "total_databases": len(self.db_manager.list_databases()),
            "initialized_databases": len(self.db_manager._initialized_databases),
            "databases_by_role": {},
            "databases_by_type": {}
        }
        
        # 按角色分组
        for role in DatabaseRole:
            databases = self.db_manager.get_databases_by_role(role)
            if databases:
                status["databases_by_role"][role.value] = databases
        
        # 按类型分组
        for db_type in DatabaseType:
            databases = self.db_manager.get_databases_by_type(db_type)
            if databases:
                status["databases_by_type"][db_type.value] = databases
        
        return status


class MockAsyncSession:
    """模拟异步会话（用于测试或无SQLAlchemy环境）"""
    
    async def execute(self, query: str, params: Optional[Dict] = None):
        """模拟执行查询"""
        logger.info(f"模拟执行: {query}")
        return MockQueryResult()
    
    async def commit(self):
        """模拟提交"""
        logger.info("模拟提交事务")
    
    async def rollback(self):
        """模拟回滚"""
        logger.info("模拟回滚事务")
    
    async def close(self):
        """模拟关闭"""
        logger.info("模拟关闭会话")
    
    async def begin(self):
        """模拟开始事务"""
        logger.info("模拟开始事务")
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()


class MockQueryResult:
    """模拟查询结果"""
    
    def fetchall(self):
        return []
    
    def fetchone(self):
        return None
    
    def scalar(self):
        return None


# 便利函数
def create_database_manager(databases_config: Dict[str, Dict[str, Any]]) -> DatabaseManager:
    """
    创建数据库管理器的便利函数
    
    Args:
        databases_config: 命名数据库配置字典
        
    Returns:
        DatabaseManager实例
    """
    return DatabaseManager(databases_config)


def create_database_router(databases_config: Dict[str, Dict[str, Any]]) -> DatabaseRouter:
    """
    创建数据库路由器的便利函数
    
    Args:
        databases_config: 命名数据库配置字典
        
    Returns:
        DatabaseRouter实例
    """
    db_manager = DatabaseManager(databases_config)
    return DatabaseRouter(db_manager)
